package org.gwh.airagknowledge.core.embedding;

import com.pgvector.PGvector;
import lombok.extern.slf4j.Slf4j;
import org.gwh.airagknowledge.entity.DocumentChunk;
import org.gwh.airagknowledge.entity.VectorEmbedding;
import org.gwh.airagknowledge.repository.VectorEmbeddingRepository;
import org.springframework.context.annotation.Primary;
import org.springframework.stereotype.Component;
import org.springframework.transaction.annotation.Transactional;

import java.util.List;
import java.util.Random;

/**
 * EmbeddingGenerator的测试替代类
 * 生成随机的嵌入向量，而不是使用真实的嵌入模型
 */
@Slf4j
@Component
@Primary
public class MockEmbeddingGenerator {

    private final VectorEmbeddingRepository vectorEmbeddingRepository;
    private final Random random = new Random(42); // 固定种子以保证可重复性

    public MockEmbeddingGenerator(VectorEmbeddingRepository vectorEmbeddingRepository) {
        this.vectorEmbeddingRepository = vectorEmbeddingRepository;
    }

    /**
     * 生成随机的嵌入向量
     * @param chunk 文档块
     * @return 向量嵌入对象
     */
    @Transactional
    public VectorEmbedding generateEmbedding(DocumentChunk chunk) {
        try {
            // 生成随机嵌入向量 (384维，与AllMiniLmL6V2EmbeddingModel一致)
            float[] vector = new float[384];
            for (int i = 0; i < vector.length; i++) {
                vector[i] = random.nextFloat() * 2 - 1; // 在 -1 到 1 之间的随机值
            }
            
            // 创建PGvector对象
            PGvector pgVector = new PGvector(vector);
            
            // 创建并保存向量嵌入
            VectorEmbedding embedding = VectorEmbedding.builder()
                    .chunk(chunk)
                    .embedding(pgVector.getVector())
                    .build();
            
            return vectorEmbeddingRepository.save(embedding);
        } catch (Exception e) {
            log.error("Error generating mock embedding for chunk ID: {}", chunk.getId(), e);
            throw new RuntimeException("Failed to generate mock embedding", e);
        }
    }

    /**
     * 生成随机的查询嵌入向量
     * @param query 查询字符串
     * @return 嵌入向量
     */
    public float[] generateQueryEmbedding(String query) {
        float[] vector = new float[384];
        for (int i = 0; i < vector.length; i++) {
            vector[i] = random.nextFloat() * 2 - 1; // 在 -1 到 1 之间的随机值
        }
        return vector;
    }

    /**
     * 生成随机的查询PGVector
     * @param query 查询字符串
     * @return PGVector对象
     */
    public PGvector generateQueryPGVector(String query) {
        return new PGvector(generateQueryEmbedding(query));
    }
    
    /**
     * 查找相似的嵌入
     * @param knowledgeBaseId 知识库ID
     * @param query 查询字符串
     * @param limit 限制数量
     * @return 相似的向量嵌入列表
     */
    public List<VectorEmbedding> findSimilarEmbeddings(Long knowledgeBaseId, String query, int limit) {
        PGvector queryVector = generateQueryPGVector(query);
        return vectorEmbeddingRepository.findSimilarEmbeddings(knowledgeBaseId, queryVector, limit);
    }
} 